package vllm

import (
	"bytes"
	"encoding/json"
	"fmt"
	"golang.org/x/net/context"
	"io"
	"net/http"
	"net/url"
)

type Predict struct {
	Completion string
}

type VLLMClient struct {
	endpoint  string
	modelPath string
}

func NewVLLMClient(endpoint, modelPath string) *VLLMClient {
	return &VLLMClient{
		endpoint:  endpoint,
		modelPath: modelPath,
	}
}

type vLLMRequest struct {
	Model        string    `json:"model"`
	Messages     []Message `json:"messages"`
	MaxTokens    int       `json:"max_tokens"`
	Temperature  float64   `json:"temperature"`
	TopP         float64   `json:"top_p"`
	N            int       `json:"n,omitempty"`
	Stream       bool      `json:"stream"`
	ChatTemplate string    `json:"chat_template,omitempty"`
}

type vLLMResponse struct {
	Choices []struct {
		Message Message `json:"message"`
	} `json:"choices"`
}

func (v *VLLMClient) Predict(ctx context.Context, opts ...PredictOption) (*Predict, error) {
	po := NewPredictOptions(opts...)
	vllmReq := vLLMRequest{
		Model:        po.Model,
		Messages:     po.Messages,
		MaxTokens:    po.MaxTokens,
		Temperature:  po.Temperature,
		TopP:         po.TopP,
		Stream:       po.Stream,
		ChatTemplate: po.ChatTemplate,
	}
	body, err := json.Marshal(vllmReq)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}
	apiURL := "/v1/chat/completions"
	urlPath, err := url.JoinPath(v.endpoint, apiURL)
	if err != nil {
		return nil, fmt.Errorf("failed to join path: %w", err)
	}
	reqHTTP, err := http.NewRequestWithContext(ctx, "POST", urlPath, bytes.NewBuffer(body))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}
	reqHTTP.Header.Set("Content-Type", "application/json")
	client := http.DefaultClient
	resp, err := client.Do(reqHTTP)
	if err != nil {
		return nil, fmt.Errorf("failed to call vLLM: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		bodyBytes, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("vLLM returned status code %d: %s", resp.StatusCode, bodyBytes)
	}
	var result vLLMResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("failed to decode vLLM response: %w", err)
	}
	if len(result.Choices) == 0 {
		return nil, fmt.Errorf("no choices in vLLM response")
	}
	return &Predict{Completion: result.Choices[0].Message.Content}, nil
}
